# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import functools
import base64
import json
import hashlib
import memcache

KEY_MAX_LENGTH = 250


@functools.lru_cache(maxsize=2048)
def to_well_formed_key(key):
  if isinstance(key, tuple):
    key = json.dumps(key)
  assert isinstance(key, str), key
  valid_key = base64.b64encode(key.encode())
  key_len = len(valid_key)
  if key_len > KEY_MAX_LENGTH:
    valid_key = hashlib.md5(valid_key).hexdigest().upper()
  else:
    valid_key = valid_key.decode()
  return valid_key


class MemcachedConnector(object):
  def __init__(self, servers=None):
    self._servers = servers or ['127.0.0.1:11211']
    self._client = memcache.Client(self._servers)
    assert self._client.get_stats(), 'Memcached service not on.'

  def get(self, key):
    key = to_well_formed_key(key)
    value = self._client.get(key)
    return value

  def set(self, key, value):
    key = to_well_formed_key(key)
    self._client.set(key, json.dumps(value))
